package org.mm.proxycache;

import java.net.InetSocketAddress;
import java.util.Queue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ConcurrentMap;

import org.apache.mina.common.ConnectFuture;
import org.apache.mina.common.IoConnector;
import org.apache.mina.common.IoFutureListener;
import org.apache.mina.common.IoSession;
import org.apache.mina.filter.codec.ProtocolCodecFilter;
import org.apache.mina.protocol.dns.DnsFuture;
import org.apache.mina.protocol.dns.DnsResolver;
import org.apache.mina.transport.socket.nio.NioSocketConnector;
import org.mm.proxycache.http.codec.HttpCodecFactory;
import org.mm.proxycache.util.SessionDataStructureFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.inject.Inject;
import com.google.inject.Provider;
import com.google.inject.Singleton;

@Singleton
public class ConnectionManager
{
    static Logger log = LoggerFactory.getLogger( ConnectionManager.class );

    private static final int CONNECT_TIMEOUT = 15;
    private static final int IDLE_TIMEOUT = 30;

    private final DnsResolver dnsResolver = new DnsResolver();

    private final ConcurrentMap<String, Queue<IoSession>> connections = new ConcurrentHashMap<String, Queue<IoSession>>();

    private final IoConnector connector;

    @Inject
    public ConnectionManager( Provider<ProxyServerHandler> handlerProvider,
            Provider<ProxyClientHttpsFilterBypass> proxyClientHttpsFilterProvider )
    {
        // Create a new connector
        connector = new NioSocketConnector( 1 );
        connector.setConnectTimeout( CONNECT_TIMEOUT );

        connector.getFilterChain().addLast( "https-bypasser", proxyClientHttpsFilterProvider.get() );

        connector.getFilterChain().addLast( "http-codec",
                new ProtocolCodecFilter( new HttpCodecFactory() ) );
        connector.setHandler( handlerProvider.get() );
        connector.setSessionDataStructureFactory( new SessionDataStructureFactory() );
    }

    public void getConnection( final String host, final int port, final IoSession clientSession,
            final ConnectionCallback callback )
    {
        log.debug( "Client Session: {}", clientSession );
        final String key = String.format( "%s:%d", host, port );

        // If already open, reuse the connection
        Queue<IoSession> sessionList = connections.get( key );
        if ( sessionList != null )
        {
            while ( true )
            {
                IoSession session = sessionList.poll();
                if ( session == null )
                {
                    // No connection for this host left in cache
                    connections.remove( key );
                    break;
                }

                if ( !session.isClosing() )
                {
                    // Found a cached opened connection
                    session.setAttribute( ClientSessionName, clientSession );
                    session.getConfig().setWriterIdleTime( 0 );
                    callback.done( session );
                    return;
                }
            }
        }
        
        DnsFuture dnsFuture = dnsResolver.resolve( host );
        dnsFuture.addListener( new IoFutureListener<DnsFuture>() {

            public void operationComplete( DnsFuture dnsFuture )
            {
                if ( dnsFuture.isNameError() )
                {
                    log.error( "Error resolving hostname: {}", host );
                    callback.error();
                    return;
                }
                
                ConnectFuture future = connector.connect( new InetSocketAddress( dnsFuture.get(), port ) );
                future.addListener( new IoFutureListener<ConnectFuture>() {

                    public void operationComplete( ConnectFuture future )
                    {
                        if ( future.isConnected() )
                        {
                            IoSession session = future.getSession();
                            session.getConfig().setBothIdleTime( 0 );
                            session.setAttribute( ClientSessionName, clientSession );
                            session.setAttribute( ConnectionEndpoint, key );

                            callback.done( session );
                        }
                        else
                        {
                            log.error( "Error connecting to server" );
                            callback.error();
                        }
                    }

                } );
            }
            
        });
    }

    public void release( IoSession session )
    {
        String key = (String) session.getAttribute( ConnectionEndpoint );
        Queue<IoSession> queue = new ConcurrentLinkedQueue<IoSession>();
        Queue<IoSession> newQueue = connections.putIfAbsent( key, queue );
        if ( newQueue != null )
            queue = newQueue;

        if ( session.getConfig().getBothIdleTime() != 0 )
            // Modify the timeout only if wasn't set by server
            session.getConfig().setBothIdleTime( IDLE_TIMEOUT );

        queue.add( session );
    }

    public interface ConnectionCallback
    {
        void done( IoSession session );

        void error();
    }

    private static final String ConnectionEndpoint = "ConnectionManager-Endpoint";
    static final String ClientSessionName = "ConnectionManager-ClientSession";

}
